import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import akshare as ak
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error
import warnings
warnings.filterwarnings('ignore')
plt.rcParams["font.sans-serif"] = ["Songti SC"]  # 设置中文字体（Windows）
plt.rcParams["axes.unicode_minus"] = False    # 解决负号显示异常
class StockAnalyzer:
    def __init__(self, stock_code, start_date, end_date):
        """
        初始化股票分析器
        :param stock_code: 股票代码，如 'sh000001' 或 'sz000001'
        :param start_date: 开始日期，格式 '20200101'
        :param end_date: 结束日期，格式 '20230101'
        """
        self.stock_code = stock_code
        self.start_date = start_date
        self.end_date = end_date
        self.data = None
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        
    def fetch_data(self):
        """获取股票数据"""
        try:
            # 使用akshare获取股票日线数据
            self.data = ak.stock_zh_a_hist(symbol=self.stock_code[2:], period="daily", 
                                         start_date=self.start_date, end_date=self.end_date, adjust="qfq")
            # print(self.data.columns)
            # 重命名列名以便于处理
            self.data.columns = ['日期','股票代码', '开盘', '收盘', '最高', '最低', '成交量', '成交额', '振幅', '涨跌幅', '涨跌额', '换手率']
            print(f"成功获取{self.stock_code}的股票数据，共{len(self.data)}行记录")
        except Exception as e:
            print(f"获取股票数据失败: {e}")
            return False
        return True
    
    def preprocess_data(self):
        """数据预处理"""
        if self.data is None or self.data.empty:
            print("没有数据可处理")
            return False
        
        # 确保数据类型正确
        numeric_columns = ['开盘', '收盘', '最高', '最低', '成交量', '成交额', '涨跌幅', '涨跌额', '换手率']
        for col in numeric_columns:
            self.data[col] = pd.to_numeric(self.data[col], errors='coerce')
        
        # 处理缺失值
        self.data.fillna(method='ffill', inplace=True)
        
        # 计算技术指标
        self.calculate_technical_indicators()
        
        # 移除包含NaN的行
        self.data.dropna(inplace=True)
        
        # 重置索引
        self.data.reset_index(drop=True, inplace=True)
        
        print("数据预处理完成")
        return True
    
    def calculate_technical_indicators(self):
        """计算技术指标"""
        # 移动平均线
        self.data['MA5'] = self.data['收盘'].rolling(window=5).mean()
        self.data['MA10'] = self.data['收盘'].rolling(window=10).mean()
        self.data['MA20'] = self.data['收盘'].rolling(window=20).mean()
        
        # 相对强弱指数 (RSI)
        delta = self.data['收盘'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        self.data['RSI'] = 100 - (100 / (1 + rs))
        
        # 移动平均收敛散度 (MACD)
        self.data['EMA12'] = self.data['收盘'].ewm(span=12, adjust=False).mean()
        self.data['EMA26'] = self.data['收盘'].ewm(span=26, adjust=False).mean()
        self.data['MACD'] = self.data['EMA12'] - self.data['EMA26']
        self.data['Signal'] = self.data['MACD'].ewm(span=9, adjust=False).mean()
        self.data['Histogram'] = self.data['MACD'] - self.data['Signal']
        
        # 布林带
        self.data['STD20'] = self.data['收盘'].rolling(window=20).std()
        self.data['UpperBand'] = self.data['MA20'] + (self.data['STD20'] * 2)
        self.data['LowerBand'] = self.data['MA20'] - (self.data['STD20'] * 2)
        
        # 成交量变化率
        self.data['VolumeChange'] = self.data['成交量'].pct_change() * 100
        
        # 价格变化率
        self.data['PriceChange'] = self.data['收盘'].pct_change() * 100
        
    def prepare_for_modeling(self, target_column='收盘', days_ahead=5):
        """
        准备建模数据
        :param target_column: 目标列
        :param days_ahead: 预测未来天数
        """
        if self.data is None or self.data.empty:
            print("没有数据可处理")
            return None, None, None, None, None
        
        # 创建目标变量 - 未来n天的收盘价
        self.data[f'未来{days_ahead}天收盘'] = self.data[target_column].shift(-days_ahead)
        
        # 移除包含NaN的行
        model_data = self.data.dropna().copy()
        
        # 特征选择
        features = [
            '开盘', '收盘', '最高', '最低', '成交量', '成交额', 
            'MA5', 'MA10', 'MA20', 'RSI', 'MACD', 'Signal', 'Histogram',
            'UpperBand', 'LowerBand', 'VolumeChange', 'PriceChange'
        ]
        
        X = model_data[features]
        y = model_data[f'未来{days_ahead}天收盘']
        
        # 数据缩放
        X_scaled = self.scaler.fit_transform(X)
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y, test_size=0.2, shuffle=False
        )
        
        # 保存用于预测的最近数据
        recent_data = self.data[features].tail(days_ahead)
        recent_data_scaled = self.scaler.transform(recent_data)
        
        return X_train, X_test, y_train, y_test, recent_data_scaled
    
    def train_model(self, X_train, y_train):
        """训练随机森林回归模型"""
        model = RandomForestRegressor(
            n_estimators=100, 
            max_depth=10, 
            random_state=42,
            n_jobs=-1
        )
        model.fit(X_train, y_train)
        return model
    
    def evaluate_model(self, model, X_test, y_test):
        """评估模型性能"""
        y_pred = model.predict(X_test)
        
        # 计算评估指标
        mse = mean_squared_error(y_test, y_pred)
        rmse = np.sqrt(mse)
        mae = mean_absolute_error(y_test, y_pred)
        
        print(f"模型评估结果:")
        print(f"均方误差 (MSE): {mse:.4f}")
        print(f"均方根误差 (RMSE): {rmse:.4f}")
        print(f"平均绝对误差 (MAE): {mae:.4f}")
        
        return {
            'mse': mse,
            'rmse': rmse,
            'mae': mae,
            'actual': y_test.values,
            'predicted': y_pred
        }
    
    def predict_future(self, model, recent_data_scaled, days_ahead=5):
        """预测未来走势"""
        future_predictions = model.predict(recent_data_scaled)
        return future_predictions
    
    def plot_results(self, model_data, y_test, y_pred, future_predictions, days_ahead=5):
        """可视化分析结果"""
        plt.figure(figsize=(16, 12))
        
        # 绘制历史价格
        plt.subplot(2, 2, 1)
        plt.plot(model_data['日期'], model_data['收盘'], label='收盘价')
        plt.plot(model_data['日期'], model_data['MA5'], label='5日均线')
        plt.plot(model_data['日期'], model_data['MA10'], label='10日均线')
        plt.plot(model_data['日期'], model_data['MA20'], label='20日均线')
        plt.title('历史股价与均线')
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True)
        
        # 绘制RSI指标
        plt.subplot(2, 2, 2)
        plt.plot(model_data['日期'], model_data['RSI'])
        plt.axhline(y=70, color='r', linestyle='-', label='超买')
        plt.axhline(y=30, color='g', linestyle='-', label='超卖')
        plt.title('相对强弱指数 (RSI)')
        plt.xlabel('日期')
        plt.ylabel('RSI值')
        plt.legend()
        plt.grid(True)
        
        # 绘制MACD指标
        plt.subplot(2, 2, 3)
        plt.plot(model_data['日期'], model_data['MACD'], label='MACD')
        plt.plot(model_data['日期'], model_data['Signal'], label='Signal')
        plt.bar(model_data['日期'], model_data['Histogram'], label='Histogram')
        plt.title('MACD指标')
        plt.xlabel('日期')
        plt.legend()
        plt.grid(True)
        
        # 绘制预测结果
        test_dates = model_data['日期'].iloc[-len(y_test):]
        future_dates = pd.date_range(start=test_dates.iloc[-1], periods=days_ahead+1, freq='B')[1:]
        
        plt.subplot(2, 2, 4)
        plt.plot(test_dates, y_test, label='实际价格')
        plt.plot(test_dates, y_pred, label='预测价格', alpha=0.7)
        
        # 预测未来价格
        plt.plot(future_dates, future_predictions, 'g--', label='未来预测', alpha=0.7)
        
        plt.title('股价预测结果')
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True)
        
        plt.tight_layout()
        plt.show()


def main():
    # 设置参数
    stock_code = 'sh002235'  # 上证指数
    start_date = '20240601'
    end_date = '20250625'
    days_ahead = 5  # 预测未来5天
    
    # 创建分析器实例
    analyzer = StockAnalyzer(stock_code, start_date, end_date)
    
    # 获取数据
    if not analyzer.fetch_data():
        return
    
    # 数据预处理
    if not analyzer.preprocess_data():
        return
    
    # 准备建模数据
    X_train, X_test, y_train, y_test, recent_data = analyzer.prepare_for_modeling(
        target_column='收盘', 
        days_ahead=days_ahead
    )
    
    if X_train is None:
        return
    
    # 训练模型
    model = analyzer.train_model(X_train, y_train)
    
    # 评估模型
    evaluation = analyzer.evaluate_model(model, X_test, y_test)
    
    # 预测未来
    future_predictions = analyzer.predict_future(model, recent_data, days_ahead)
    
    print(f"\n未来{days_ahead}天的股价预测:")
    for i, date in enumerate(pd.date_range(start=pd.to_datetime(end_date) + pd.Timedelta(days=1), periods=days_ahead, freq='B')):
        print(f"{date.strftime('%Y-%m-%d')}: {future_predictions[i]:.2f}")
    
    # 可视化结果
    model_data = analyzer.data.dropna()
    analyzer.plot_results(model_data, evaluation['actual'], evaluation['predicted'], future_predictions, days_ahead)


if __name__ == "__main__":
    main()    
